-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Add metadata to the Agent class. #3370
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
| def _run_span_end_attributes( | ||
| self, | ||
| settings: InstrumentationSettings, | ||
| usage: _usage.RunUsage, | ||
| message_history: list[_messages.ModelMessage], | ||
| new_message_index: int, | ||
| graph_ctx: GraphRunContext[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, OutputDataT]], | ||
| ): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I reworked what attributes are passed to this because the existing ones are already available in the RunContext. I hope that's okay.
|
@sstanovnik Have you seen Baggage https://logfire.pydantic.dev/docs/reference/advanced/baggage/#basic-usage That could also be a solution here that doesn't require us to add any new fields. |
docs/logfire.md
Outdated
| from pydantic_ai.models.instrumented import InstrumentationSettings | ||
|
|
||
|
|
||
| def span_attribute_callback(ctx: RunContext[None]) -> dict[str, str]: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is useful, but not type safe unfortunately, as InstrumentationSettings is not generic in the deps type (and probably shouldn't be), meaning that you could create a span_attribute_callback that takes RunContext[Foo] and use it on an Agent[Bar] without any typing issues, which would then fail at runtime if you try to read attrs that exist on Foo but not Bar...
The obvious ways to solve that would be to make InstrumentationSettings generic or to make Agent and agent.run take the span_attributes directly, but I don't like either of those :) If baggage works for you, I'd be inclined to not add any new feature to Pydantic AI and just document that instead.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, I noticed that this couldn't be made type-safe without making the carrier generic, but that seems like a huge undertaking that also changes the interface in a way that, I would think, isn't desirable until v2.
Baggage would work for my immediate usecase, even though it applies to all child spans, which I don't really want it to - conceptually, I'm wanting to add data to the one specific span.
What about only accepting a dict for now, and leaving a callable taking RunContext[TDeps] for whenever InstrumentationSettings can be adapted to take a generic parameter?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@sstanovnik Maybe instrumentation settings are not the best place for this:
In line with #3263, what do you think about arbitrary metadata: dict[str, Any] on the Agent, that's also set on the agent run span under a metadata attribute? That way you wouldn't be setting span attributes directly (assuming that's not an issue), but they would be queryable etc, and it'd feel more natural to add this on both Agent and agent.run directly, with support for RunContext-taking callables.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That's also great - any way of adding information to the span would be great. What do you mean by
That way you wouldn't be setting span attributes directly (assuming that's not an issue), but they would be queryable [...]
I see that the linked PR does add the metadata into the span.
I'm up for closing this PR and opening another with the same approach using metadata, if you'd like :)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@sstanovnik I just mean that the metadata dict keys will not directly become span attributes, as they'd be nested under metadata, but if that's fine for your use case, feel free to update this PR (with force push possibly), or create a new one!
Metadata is attached to the logfire.agent.metadata span attribute. It can either be a string, dict, or a callable taking the RunContext and returning a string or a dict.
c57e1bb to
11687ce
Compare
sstanovnik
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Force-pushed this to retain the relevant conversation history of the previous implementation. I also added some comments on particulars I'm not certain about.
A major difference from your suggestion is that I did not add this to the Agent.run method. I felt the change was large enough as it is without having to figure out how to wire/override/merge that.
| }, | ||
| ) | ||
|
|
||
| run_metadata: str | dict[str, str] | None = None |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We could technically set a non-callable metadata even if the run fails, and the ctx isn't available. What do you think? I opted for the simpler implementation.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If the run fails, wouldn't we still have RunContext, or at least deps? I think we should try to always store the metadata.
We should also document explicitly whether the callable is executed at the start or end of the run, as RunContext would change (e.g. messages, usage); I think at the end is best.
docs/logfire.md
Outdated
| ### Adding Custom Metadata | ||
|
|
||
| Use the agent's `metadata` parameter to attach additional data to the agent's span. | ||
| Metadata can be provided as a string, a dictionary, or a callable that reads the [`RunContext`][pydantic_ai.tools.RunContext] to compute values on each run. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'd prefer to only support dict[str, Any] (or a callable returning that), as we do for other metadata fields.
docs/logfire.md
Outdated
| ) | ||
| ``` | ||
|
|
||
| When instrumentation is enabled, the resolved metadata is recorded on the agent span under the `logfire.agent.metadata` attribute. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I may still change my mind on this, but I don't think it needs the logfire. prefix as it's user-provided rather than Logfire specific (like some of the Evals stuff is)
|
|
||
| ### Adding Custom Metadata | ||
|
|
||
| Use the agent's `metadata` parameter to attach additional data to the agent's span. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Right now this is very instrumentation specific, although the name intentionally isn't. I think we should also make the resulting metadata available on AgentRunResult (and the other classes that have a run_id field) so the user can read/store it after the fact.
| }, | ||
| ) | ||
|
|
||
| run_metadata: str | dict[str, str] | None = None |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If the run fails, wouldn't we still have RunContext, or at least deps? I think we should try to always store the metadata.
We should also document explicitly whether the callable is executed at the start or end of the run, as RunContext would change (e.g. messages, usage); I think at the end is best.
| toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET, | ||
| tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] | _utils.Unset = _utils.UNSET, | ||
| instructions: Instructions[AgentDepsT] | _utils.Unset = _utils.UNSET, | ||
| metadata: AgentMetadataValue[AgentDepsT] | _utils.Unset = _utils.UNSET, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If we're doing override, I'd like to do run (and all the derived methods...) as well. That run arg would be merged into the agent-construction-time metadata.
|
This PR is stale, and will be closed in 3 days if no reply is received. |
|
Apologies for the delay in responding to your review. Thank you for taking the time for this. I believe I've addressed all comments:
|
| You can retrieve usage statistics (tokens, requests, etc.) at any time from the [`AgentRun`][pydantic_ai.agent.AgentRun] object via `agent_run.usage()`. This method returns a [`RunUsage`][pydantic_ai.usage.RunUsage] object containing the usage data. | ||
|
|
||
| Once the run finishes, `agent_run.result` becomes a [`AgentRunResult`][pydantic_ai.agent.AgentRunResult] object containing the final output (and related metadata). | ||
| You can inspect [`agent_run.metadata`][pydantic_ai.agent.AgentRun] or [`agent_run.result.metadata`][pydantic_ai.agent.AgentRunResult] after the run completes to read any metadata configured on the agent. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This deserves a separate section that also explains how to set it. I think much of what's currently in logfire.md can be moved here, and then those docs can link here.
| 'openai:gpt-5', | ||
| instrument=True, | ||
| metadata=lambda ctx: {'deployment': 'staging', 'tenant': ctx.deps.tenant}, | ||
| ) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It'd be nice to expand the example to show running the agent with some deps.tenant, and then printing result.metadata at the end
| model_settings: ModelSettings | None = None, | ||
| usage_limits: _usage.UsageLimits | None = None, | ||
| usage: _usage.RunUsage | None = None, | ||
| metadata: AgentMetadataValue[AgentDepsT] | None = None, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's call it AgentMetadata, without the Value
| usage: Optional usage to start with, useful for resuming a conversation or agents used in tools. | ||
| metadata: Optional metadata to attach to this run. Accepts a dictionary or a callable taking | ||
| [`RunContext`][pydantic_ai.tools.RunContext]. The resolved dictionary is shallow merged into the | ||
| agent's metadata (or any [`Agent.override`][pydantic_ai.agent.Agent.override]) with run-level keys |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Existing agent.overrides always cause the agent.run value to be ignored, fully overwriting the agent + agent run original. We should have the same behavior here
| try: | ||
| yield agent_run | ||
| finally: | ||
| resolve_run_metadata() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does this need to be a function or could it be inlined?
| finally: | ||
| run_span.end() | ||
|
|
||
| def _compute_agent_metadata( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For consistency with the other fields supported by agent.override, let's call this _get_metadata() and the run_metadata_config additional_metadata
| base_config = self._metadata | ||
|
|
||
| base_metadata = self._resolve_metadata_config(base_config, ctx) | ||
| run_metadata = self._resolve_metadata_config(run_metadata_config, ctx) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As mentioned above, overridden metadata fully overrides run metadata
| attrs['pydantic_ai.variable_instructions'] = True | ||
|
|
||
| if metadata is not None: | ||
| attrs['agent.metadata'] = json.dumps(metadata) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's just make this metadata like we do with evals:
pydantic-ai/pydantic_evals/pydantic_evals/dataset.py
Lines 300 to 301 in c27e5e4
| if metadata is not None: | |
| extra_attributes['metadata'] = metadata |
| usage: Optional usage to start with, useful for resuming a conversation or agents used in tools. | ||
| metadata: Optional metadata to attach to this run. Accepts a dictionary or a callable taking | ||
| [`RunContext`][pydantic_ai.tools.RunContext]. The resolved dictionary is shallow merged into the | ||
| agent's metadata (or any [`Agent.override`][pydantic_ai.agent.Agent.override]) with run-level keys |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
See above, this behavior is inconsistent with existing expectations, so the behavior and docstring will need updating
| def metadata(self) -> dict[str, Any] | None: | ||
| """Metadata associated with this agent run, if configured.""" | ||
| if self._metadata_getter is not None: | ||
| return self._metadata_getter() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I understand that we do this because the metadata won't be available on self._run_ctx as it's only built at the end of the run, but I wonder if we can change that and simplify this.
Can we make sure the metadata is always available, just initially built off the original RunContext and then regenerated at the end of the run? Usually it'll just use deps anyway, which is available from the start.
This allows users to add extra metadata to the agent's span, either as a direct string or dict, or a callable that takes the RunContext. The attributes are added/computed after the agent finishes. The metadata is under the
logfire.agent.metadataattribute.